import numpy as np
import pandas as pd

from datetime import datetime

# Visualization libraries
import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt
import plotly.express as px

# ML/DL libraries
import tensorflow as tf
from tensorflow.keras import layers
import tensorflow.keras.applications as models
import tensorflow_ranking as tfr
import tensorflow_recommenders as tfrs
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
import keras_tuner as kt
from scipy.stats import kendalltau, spearmanr

# Misc

import math
import random
from joblib import Parallel, delayed
import PIL
from PIL import Image
import os
import pathlib
import warnings
from pprint import pprint
from numba import cuda 


def dcg(scores):
    scores_ = np.array(scores)
    discounts = np.log2(np.arange(1, len(scores_) + 1) + 1)
    return np.sum(scores_ / discounts)


def ndcg(y_true_batch, y_pred_batch):
    y_true_batch_ = np.array(y_true_batch)
    y_pred_batch_ = np.array(y_pred_batch)
    ndcg_scores = []
    
    for y_true, y_pred in zip(y_true_batch_, y_pred_batch_):
        y_true = np.array(y_true)
        y_pred = np.array(y_pred)

        order = np.argsort(-y_true)
        y_pred_sorted_by_true = y_pred[order]

        y_true_ideal = np.sort(y_true)[::-1]
        idcg = dcg(y_true_ideal)

        dcg_score = dcg(y_pred_sorted_by_true)
        
        ndcg_score = dcg_score / idcg if idcg > 0 else 0.0
        ndcg_scores.append(ndcg_score)
    
    return np.mean(ndcg_scores)


def permutation_accuracy(y_true, y_pred):
    y_true_ = np.array(y_true)
    y_pred_ = np.array(y_pred)
    
    correct = tf.equal(y_true_, y_pred_)
    
    return tf.reduce_mean(tf.cast(correct, tf.float32))


def accuracy_by_rank(y_true_batch, y_pred_batch, rank):
    y_true_batch_ = tf.convert_to_tensor(y_true_batch)
    y_pred_batch_ = tf.convert_to_tensor(y_pred_batch)
    
    relevant_true_positions = tf.cast(tf.equal(y_true_batch_, rank), tf.float32)
    relevant_pred_positions = tf.cast(tf.equal(y_pred_batch_, rank), tf.float32)
    
    correct_predictions = tf.multiply(relevant_true_positions, relevant_pred_positions)
    
    correct_predictions_sum = tf.reduce_sum(correct_predictions, axis=1)
    relevant_true_positions_sum = tf.reduce_sum(relevant_true_positions, axis=1)
    
    accuracy = tf.where(relevant_true_positions_sum > 0, 
                        correct_predictions_sum / relevant_true_positions_sum, 
                        tf.zeros_like(correct_predictions_sum))
    
    mean_accuracy = tf.reduce_mean(accuracy)
    
    return mean_accuracy.numpy()


def accuracy_at_k(y_true_batch, y_pred_batch):
    y_true_batch_ = np.array(y_true_batch)
    y_pred_batch_ = np.array(y_pred_batch)

    accuracies = []

    for y_true, y_pred in zip(y_true_batch_, y_pred_batch_):
        y_true = np.array(y_true)
        y_pred = np.array(y_pred)

        order = np.argsort(-y_true)
        y_pred_sorted_by_true = y_pred[order]
        y_true_ideal = np.sort(y_true)[::-1]
        
        cur_accuracies = []
    
        for k in range(5):
            y_pred_top_k = y_pred_sorted_by_true[:k]
            y_true_top_k = y_true_ideal[:k]
            
            correct_count = np.isin(y_pred_top_k, y_true_top_k).sum()
            
            accuracy = correct_count / k
            cur_accuracies.append(accuracy)
        accuracies.append(np.array(cur_accuracies))
    
    return np.mean(accuracies, axis=0)


def batch_kendalls_tau(y_true_batches, y_pred_batches):
    tau_values = []
    for y_true, y_pred in zip(y_true_batches, y_pred_batches):
        y_true = np.array(y_true)
        y_pred = np.array(y_pred)

        order = np.argsort(-y_true)
        y_pred_sorted_by_true = y_pred[order]
        y_true_ideal = np.sort(y_true)[::-1]

        tau, _ = kendalltau(y_true_ideal, y_pred_sorted_by_true)
        tau_values.append(tau)
    return np.mean(tau_values)


def batch_spearmans_rho(y_true_batches, y_pred_batches):
    rho_values = []
    for y_true, y_pred in zip(y_true_batches, y_pred_batches):
        y_true = np.array(y_true)
        y_pred = np.array(y_pred)

        order = np.argsort(-y_true)
        y_pred_sorted_by_true = y_pred[order]
        y_true_ideal = np.sort(y_true)[::-1]

        rho, _ = spearmanr(y_true_ideal, y_pred_sorted_by_true)
        rho_values.append(rho)
    return np.mean(rho_values)


def ordered_pair_accuracy(y_true_batches, y_pred_batches):
    def single_batch_opa(y_true_, y_pred_):
        y_true = np.array(y_true_)
        y_pred = np.array(y_pred_)

        order = np.argsort(-y_true)
        y_pred_sorted_by_true = y_pred[order]
        y_true_ideal = np.sort(y_true)[::-1]

        y_true = y_true_ideal
        y_pred = y_pred_sorted_by_true

        correct_pairs = 0
        total_pairs = 0
        
        n = len(y_true)
        for i in range(n):
            for j in range(i + 1, n):
                total_pairs += 1
                if (y_true[i] < y_true[j] and y_pred[i] < y_pred[j]) or (y_true[i] > y_true[j] and y_pred[i] > y_pred[j]):
                    correct_pairs += 1
                elif (y_true[i] == y_true[j] and y_pred[i] == y_pred[j]):
                    correct_pairs += 1
        
        return correct_pairs / total_pairs if total_pairs > 0 else 0

    batch_opa_values = [single_batch_opa(y_true, y_pred) for y_true, y_pred in zip(y_true_batches, y_pred_batches)]
    return np.mean(batch_opa_values)